import numpy as np
from sklearn.neighbors import NearestNeighbors

# 模拟向量数据
np.random.seed(42)
data_vectors = np.random.rand(100, 5)  # 100个5维向量


# -------------------------------
# 图索引实现
# -------------------------------
class GraphIndex:
    def __init__(self, data_vectors):
        self.data_vectors = data_vectors
        self.graph = None

    def build_graph(self, k=5):
        """
        使用 sklearn 的NearestNeighbors为每个向量构建 k 近邻图（这里 k=5），形成一个以向量索引为节点、近邻关系为边的图结构。
        """
        nbrs = NearestNeighbors(n_neighbors=k, algorithm='auto').fit(self.data_vectors)
        distances, indices = nbrs.kneighbors(self.data_vectors)
        self.graph = {i: indices[i].tolist() for i in range(len(self.data_vectors))}

    def search(self, query_vector, top_k=5):
        """
        从图中找到与查询向量最相似的k个邻居
        通过计算查询向量与所有数据向量的欧氏距离，找到距离最近的 top_k 个向量的索引并返回
        """
        query_id = len(self.data_vectors)
        distances = [np.linalg.norm(query_vector - vec) for vec in self.data_vectors]
        nearest_indices = np.argsort(distances)[:top_k]
        return nearest_indices


# 构建图索引
graph_index = GraphIndex(data_vectors)
graph_index.build_graph()
print(type(graph_index.graph))  # <class 'dict'>
print(graph_index.graph)
# {0: [0, 51, 32, 67, 37], 1: [1, 22, 43, 93, 36], 2: [2, 67, 86, 73, 38], 3: [3, 12, 50, 21, 57], 4: [4, 34, 68, 63,
# 9], 5: [5, 9, 46, 87, 28], 6: [6, 62, 14, 77, 40], 7: [7, 47, 14, 78, 77], 8: [8, 75, 35, 94, 74], 9: [9, 5, 46,
# 65, 28], 10: [10, 54, 49, 79, 55], 11: [11, 33, 56, 63, 28], 12: [12, 50, 85, 21, 3], 13: [13, 93, 34, 27, 64],
# 14: [14, 77, 40, 7, 6], 15: [15, 84, 44, 59, 24], 16: [16, 24, 63, 56, 33], 17: [17, 76, 80, 45, 39], 18: [18, 92,
# 60, 76, 31], 19: [19, 63, 70, 98, 16], 20: [20, 60, 72, 88, 92], 21: [21, 50, 12, 85, 3], 22: [22, 1, 43, 36, 71],
# 23: [23, 79, 61, 96, 78], 24: [24, 16, 15, 70, 84], 25: [25, 42, 31, 73, 70], 26: [26, 83, 30, 4, 69], 27: [27, 64,
# 93, 31, 92], 28: [28, 46, 56, 9, 65], 29: [29, 86, 57, 21, 19], 30: [30, 77, 14, 26, 97], 31: [31, 42, 27, 18, 57],
# 32: [32, 0, 51, 37, 67], 33: [33, 56, 11, 63, 16], 34: [34, 4, 56, 13, 33], 35: [35, 80, 8, 17, 75], 36: [36, 22,
# 87, 12, 71], 37: [37, 51, 53, 65, 46], 38: [38, 67, 0, 82, 80], 39: [39, 55, 76, 17, 82], 40: [40, 14, 62, 77, 6],
# 41: [41, 58, 21, 91, 3], 42: [42, 64, 31, 25, 27], 43: [43, 1, 60, 22, 92], 44: [44, 68, 15, 63, 84], 45: [45, 80,
# 17, 76, 94], 46: [46, 65, 9, 53, 37], 47: [47, 7, 81, 44, 68], 48: [48, 89, 53, 51, 37], 49: [49, 10, 54, 95, 39],
# 50: [50, 21, 12, 85, 3], 51: [51, 0, 37, 32, 53], 52: [52, 69, 79, 72, 96], 53: [53, 65, 37, 46, 85], 54: [54, 10,
# 55, 49, 39], 55: [55, 96, 72, 39, 17], 56: [56, 33, 28, 16, 11], 57: [57, 60, 31, 3, 29], 58: [58, 41, 50, 12, 21],
# 59: [59, 84, 15, 75, 37], 60: [60, 18, 76, 20, 92], 61: [61, 78, 90, 81, 7], 62: [62, 6, 40, 14, 83], 63: [63, 44,
# 16, 68, 19], 64: [64, 42, 27, 72, 31], 65: [65, 46, 53, 37, 9], 66: [66, 15, 44, 19, 98], 67: [67, 2, 38, 0, 32],
# 68: [68, 44, 98, 69, 63], 69: [69, 68, 44, 98, 52], 70: [70, 73, 86, 53, 51], 71: [71, 95, 40, 97, 83], 72: [72,
# 27, 55, 64, 20], 73: [73, 70, 2, 86, 0], 74: [74, 8, 3, 47, 68], 75: [75, 59, 8, 94, 35], 76: [76, 60, 17, 18, 39],
# 77: [77, 14, 30, 97, 40], 78: [78, 61, 7, 90, 46], 79: [79, 23, 96, 55, 52], 80: [80, 35, 17, 45, 38], 81: [81, 47,
# 61, 7, 90], 82: [82, 38, 18, 39, 55], 83: [83, 77, 40, 72, 26], 84: [84, 15, 59, 44, 24], 85: [85, 50, 12, 21, 65],
# 86: [86, 29, 2, 70, 67], 87: [87, 5, 9, 36, 85], 88: [88, 20, 18, 72, 31], 89: [89, 48, 53, 51, 37], 90: [90, 61,
# 78, 46, 81], 91: [91, 19, 4, 41, 11], 92: [92, 18, 27, 60, 43], 93: [93, 13, 27, 43, 31], 94: [94, 75, 45, 80, 8],
# 95: [95, 71, 78, 97, 40], 96: [96, 55, 37, 79, 72], 97: [97, 77, 6, 40, 14], 98: [98, 68, 69, 44, 63], 99: [99, 6,
# 62, 83, 72]}


# 图索引检索
query_vector = np.array([0.5, 0.5, 0.5, 0.5, 0.5])  # 查询向量
nearest_neighbors = graph_index.search(query_vector, top_k=3)
print("图索引中最相似的向量ID:", nearest_neighbors)
